// Copyright (c) 2022 ChenJun
// Licensed under the MIT License.

// OpenCV
#include <opencv2/core.hpp>
#include <opencv2/core/base.hpp>
#include <opencv2/core/mat.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/imgproc.hpp>

// STD
#include <cmath>
#include <vector>

#include "ThreadManager.h"
#include "armor.hpp"
#include "detector.hpp"



namespace hnurm
{

    Detector::Detector(const cv::FileNode &cfg_node)
    {
        cv::FileNode light_params_node = cfg_node["light_params"];
        cv::FileNode center_params_node = cfg_node["center_params"];
        cv::FileNode armor_params_node = cfg_node["armor_params"];
        cv::FileNode classifier_node = cfg_node["classifier_params"];

        cfg_node["min_lightness"] >> min_lightness;

        // 初始化光照参数
        light_params_node["min_ratio"] >> l.min_ratio;
        light_params_node["max_ratio"] >> l.max_ratio;
        light_params_node["max_angle"] >> l.max_angle;

        // 初始化中心R参数
        center_params_node["min_size"] >> c.min_size;
        center_params_node["max_size"] >> c.max_size;

        // 初始化装甲参数
        armor_params_node["min_light_ratio"] >> a.min_light_ratio;
        armor_params_node["min_small_center_distance"] >> a.min_small_center_distance;
        armor_params_node["max_small_center_distance"] >> a.max_small_center_distance;
        armor_params_node["min_large_center_distance"] >> a.min_large_center_distance;
        armor_params_node["max_large_center_distance"] >> a.max_large_center_distance;
        armor_params_node["max_angle"] >> a.max_angle;

        // 初始化分类器
        std::string _model_path, _label_path;
        std::vector<std::string> _ignore_classes;
        double _threshold;
        classifier_node["model_path"] >> _model_path;
        classifier_node["label_path"] >> _label_path;
        classifier_node["threshold"] >> _threshold;
        classifier_node["ignore_classes"] >> _ignore_classes;
        this->classifier = std::make_unique<NumberClassifier>(_model_path, _label_path, _threshold, _ignore_classes);
    }

    // 检测函数，输入图像，返回检测到的装甲
    std::vector<Armor> Detector::detect(const cv::Mat &input, int self_color)
    {
        // 二值化
        binary_img = preprocessImage(input);

#ifdef DEBUG_SHOW
        cv::namedWindow("preprocess",cv::WINDOW_AUTOSIZE);
        cv::imshow("preprocess", binary_img);
#endif

        // 寻找所有灯条
        lights_ = findLights(input, binary_img);

        // 灯条两两匹配
        armors_ = matchLights(lights_, self_color);

        // 裁剪数字区域的ROI图像
        if (!armors_.empty()) {
            classifier->extractNumbers(input, armors_);
            TIMEIT_ID(CLASSIFY, classifier->classify(armors_,self_color));
        }
        return armors_;
    }

    Energybuff Detector::detect_buff(const cv::Mat &input, int self_color) {
        // 二值化
        binary_img = preprocessImage(input);

#ifdef DEBUG_SHOW
        cv::namedWindow("preprocess",cv::WINDOW_AUTOSIZE);
        cv::imshow("preprocess", binary_img);
#endif
        // 寻找能量机关中心
        centers_ = findCenter(input, binary_img);
        flabellums_ = findFlabellum(input, binary_img);
        energybuff_.center = (centers_.size()==0?energybuff_.center:centers_.at(0));

        // 寻找能量机关扇叶
        
        
        if(energybuff_.flabellums.size() != 0) {
            energybuff_.init(energybuff_.flabellums.at(0));
        }
        else if(flabellums_.size() != 0) {
            energybuff_.init(flabellums_.at(0));
        }

        energybuff_.update(flabellums_);
        for(const auto&flabellum:energybuff_.flabellums) {
            if(flabellum.current_state == 1) {
                target_ = flabellum.rotate;
                energybuff_.targetarmor = flabellum;
            }
        }
        // sort(energybuff_.flabellums.begin(), energybuff_.flabellums.end(), [](Flabellum &a, Flabellum &b){return a.rate>b.rate;});
        // energybuff_.targetarmor = energybuff_.flabellums.at(0);
        return energybuff_;
    }

    // 获取所有数字图像的图像
    cv::Mat Detector::getAllNumbersImage()
    {
        if (armors_.empty()) {
            return {cv::Size(20, 28), CV_8UC1};
        } else {
            std::vector<cv::Mat> number_imgs;
            number_imgs.reserve(armors_.size());
            for (auto &armor: armors_) {
                number_imgs.emplace_back(armor.number_img);
            }
            cv::Mat all_num_img;
            cv::vconcat(number_imgs, all_num_img);
            return all_num_img;
        }
    }

    // 绘制结果函数，绘制在输入图像上
    void Detector::drawResults(cv::Mat &img)
    {
        // 绘制能量机关中心
        for (const auto &center: centers_) {
            cv::Point2f points[4];
            center.points(points);
            for (int i = 0; i < 4; i++)
            {
                line(img, points[i], points[(i + 1) % 4], cv::Scalar(0, 0, 255), 2, 8);
            }
            cv::putText(
                    img, "center_type:"+std::to_string(center.color), center.bottom, cv::FONT_HERSHEY_SIMPLEX, 0.8,
                    cv::Scalar(0, 255, 255), 2);
        }

        // 绘制能量机关扇叶
        cv::Point2f points[4];
        target_.points(points);
        for(int i=0;i<4;i++) {
            cv::line(img, points[i], points[(i + 1)%4], cv::Scalar(0, 255, 0), 2, 8);
        }
        for (const auto &flabellum:energybuff_.flabellums) {
            cv::circle(img, (flabellum.tl() + flabellum.br()) / 2, 7, cv::Scalar(0, 0, 255), -1);
            if(flabellum.current_state == 0){
                cv::rectangle(img, flabellum, cv::Scalar(226, 43, 138), 2, 8);
                cv::putText(
                    img, "id = "+std::to_string(flabellum.id)+"unlighted", flabellum.top, cv::FONT_HERSHEY_SIMPLEX, 0.8,
                    cv::Scalar(0, 255, 255), 2);
            }
            else if(flabellum.current_state == 2) {
                cv::rectangle(img, flabellum, cv::Scalar(0, 0, 255), 2, 8);
                cv::putText(
                    img, "id = "+std::to_string(flabellum.id)+"shot", flabellum.top, cv::FONT_HERSHEY_SIMPLEX, 0.8,
                    cv::Scalar(0, 255, 255), 2);
            }
            else if(flabellum.current_state == 1) {
                // cv::rectangle(img, flabellum, cv::Scalar(0, 255, 0), 2, 8);
                cv::putText(
                    img, "id = "+std::to_string(flabellum.id)+"target", flabellum.top, cv::FONT_HERSHEY_SIMPLEX, 0.8,
                    cv::Scalar(0, 255, 255), 2);
            }
        }

        // // 绘制灯条
        // for (const auto &light: lights_) {
        //     cv::circle(img, light.top, 3, cv::Scalar(255, 255, 255), 1);
        //     cv::circle(img, light.bottom, 3, cv::Scalar(255, 255, 255), 1);
        //     auto line_color = light.color == RED ? cv::Scalar(255, 255, 0) : cv::Scalar(255, 0, 255);
        //     cv::line(img, light.top, light.bottom, line_color, 1);
        //     cv::putText(
        //             img, "light_type:"+std::to_string(light.color), light.bottom, cv::FONT_HERSHEY_SIMPLEX, 0.8,
        //             cv::Scalar(0, 255, 255), 2);
        // }

        // 绘制装甲
        for (const auto &armor: armors_) {
            cv::line(img, armor.left_light.top, armor.right_light.bottom, cv::Scalar(0, 255, 0), 2);
            cv::line(img, armor.left_light.bottom, armor.right_light.top, cv::Scalar(0, 255, 0), 2);
        }

        // 显示装甲板一些参数
        for (const auto &armor: armors_) {
            cv::putText(
                    img, armor.classification_result, armor.left_light.top, cv::FONT_HERSHEY_SIMPLEX, 0.8,
                    cv::Scalar(0, 255, 255), 2);
            cv::putText(
                    img, "armor_type:"+std::to_string(armor.armor_type), cv::Point2f(armor.left_light.top.x,armor.left_light.top.y-20), cv::FONT_HERSHEY_SIMPLEX, 0.8,
                    cv::Scalar(0, 255, 255), 2);
        }
    }

    // 预处理函数，转灰度
    cv::Mat Detector::preprocessImage(const cv::Mat &rgb_img) const
    {
        // 转灰度
        cv::Mat gray_img;
        cv::cvtColor(rgb_img, gray_img, cv::COLOR_RGB2GRAY);

        // 二值化
        cv::Mat _binary_img;
        cv::threshold(gray_img, _binary_img, min_lightness, 255, cv::THRESH_BINARY);
        return _binary_img;
    }

    // 寻找灯条函数
    std::vector<Light> Detector::findLights(const cv::Mat &rbg_img, const cv::Mat &_binary_img)
    {
        using std::vector;
        vector<vector<cv::Point>> contours;
        cv::findContours(_binary_img, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

        vector<Light> lights;

        for (const auto &contour: contours) {
            if (contour.size() < 5) continue;

            auto r_rect = cv::minAreaRect(contour);
            auto light = Light(r_rect);

            if (isLight(light)) {
                auto rect = light.boundingRect();
                if (0 <= rect.x && 0 <= rect.width && rect.x + rect.width <= rbg_img.cols &&
                    0 <= rect.y && 0 <= rect.height && rect.y + rect.height <= rbg_img.rows) {
                    int sum_r = 0, sum_b = 0;
                    auto roi = rbg_img(rect);
                    // 遍历ROI
                    for (int row = 0; row < roi.rows; row++) {
                        for (int j = 0; j < roi.cols; j++) {
                            if (cv::pointPolygonTest(contour, cv::Point2f(j + rect.x, row + rect.y), false) >= 0) {
                                // 如果点在轮廓内部
                                sum_b += roi.at<cv::Vec3b>(row, j)[0];
                                sum_r += roi.at<cv::Vec3b>(row, j)[2];
                            }
                        }
                    }
                    // 红色像素的和大于蓝色像素的和？
                    float sum_br=float(sum_b)/sum_r;
                    if( sum_br > 0.99 && sum_br < 1.20){
                        light.color=PURPLE;
                        lights.emplace_back(light);
                        continue;
                    }
                    else if(sum_r>sum_b){
                        light.color =  RED;
                    }
                    else if (sum_r < sum_b){
                        light.color = BLUE;
                    }
                    lights.emplace_back(light);
                }
            }
        }

        return lights;
    }

    // 寻找能量机关中心函数
    std::vector<Center> Detector::findCenter(const cv::Mat &rbg_img, const cv::Mat &_binary_img)
    {
        using std::vector;
        cv::Mat prec = _binary_img.clone();
        vector<vector<cv::Point>> contours;
        dilate(prec, prec, cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5)));
        cv::findContours(prec, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

        vector<Center> centers;
        for (const auto &contour: contours) {
            if (contour.size() < 5) continue;

            auto r_rect = cv::minAreaRect(contour);
            auto center = Center(r_rect);
            bool iscenter = false;

            //判断是否是中心
            float ratio = center.width / center.length;
            bool ratio_ok = 0.7 < ratio && ratio < 1.3;
            bool size_ok=(center.width*center.length)>c.min_size && (center.width*center.length)<c.max_size;
            iscenter = ratio_ok && size_ok;
            //todo 添加神经网络判断

            if (iscenter) {
                auto rect = center.boundingRect();
                if (0 <= rect.x && 0 <= rect.width && rect.x + rect.width <= rbg_img.cols &&
                    0 <= rect.y && 0 <= rect.height && rect.y + rect.height <= rbg_img.rows) {
                    int sum_r = 0, sum_b = 0;
                    auto roi = rbg_img(rect);
                    // 遍历ROI
                    for (int row = 0; row < roi.rows; row++) {
                        for (int j = 0; j < roi.cols; j++) {
                            if (cv::pointPolygonTest(contour, cv::Point2f(j + rect.x, row + rect.y), false) >= 0) {
                                // 如果点在轮廓内部
                                sum_b += roi.at<cv::Vec3b>(row, j)[0];
                                sum_r += roi.at<cv::Vec3b>(row, j)[2];
                            }
                        }
                    }
                    // 红色像素的和大于蓝色像素的和？
                    float sum_br=float(sum_b)/sum_r;
                    if( sum_br > 0.99 && sum_br < 1.20){
                        center.color=PURPLE;
                        continue;
                    }
                    else if(sum_r>sum_b){
                        center.color =  RED;
                    }
                    else if (sum_r < sum_b){
                        center.color = BLUE;
                    }
                    centers.emplace_back(center);
                }
            }
        }

        return centers;
    }

    // 寻找能量机关扇叶函数
    std::vector<Flabellum> Detector::findFlabellum(const cv::Mat &rbg_img, const cv::Mat &_binary_img)
    {
        using std::vector;
        int id = 0;
        cv::Mat prec = _binary_img.clone();
        cv::Point2f center_p;
        vector<vector<cv::Point>> contours;
        dilate(prec, prec, cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5)));
        // 在内部绘制实心圆
        cv::imshow("beforedelete", prec);
        for (const auto &center: centers_) {
            center_p = center.center;
            cv::circle(prec, center_p, 50, cv::Scalar(0, 0, 0), -1);
        }
        cv::findContours(prec, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
        cv::imshow("afterdelete", prec);
        
        vector<Flabellum> flabellums;

        for (const auto &contour: contours) {
            if (contour.size() < 5) continue;

            auto r_rect = cv::boundingRect(contour);
            auto rotate = cv::minAreaRect(contour);
            auto flabellum = Flabellum(r_rect);
            bool is_flabellum = false;

            //判断扇叶
            float ratio = flabellum.width / flabellum.length;
            bool ratio_ok = 0.5 < ratio && ratio < 2 || true;
            // double size = cv::contourArea(contour);
            double size = rotate.size.height * rotate.size.width;
            bool size_ok = size > 8000 && size < 30000;
            is_flabellum = ratio_ok && size_ok;

            if (is_flabellum) {
                auto rect = flabellum;
                flabellum.id = id++;
                if (0 <= rect.x && 0 <= rect.width && rect.x + rect.width <= rbg_img.cols &&
                    0 <= rect.y && 0 <= rect.height && rect.y + rect.height <= rbg_img.rows) {
                    int sum_r = 0, sum_b = 0;
                    auto roi = rbg_img(rect);
                    // 遍历ROI
                    for (int row = 0; row < roi.rows; row++) {
                        for (int j = 0; j < roi.cols; j++) {
                            if (cv::pointPolygonTest(contour, cv::Point2f(j + rect.x, row + rect.y), false) >= 0) {
                                // 如果点在轮廓内部
                                sum_b += roi.at<cv::Vec3b>(row, j)[0];
                                sum_r += roi.at<cv::Vec3b>(row, j)[2];
                            }
                        }
                    }
                    // 红色像素的和大于蓝色像素的和？
                    float sum_br=float(sum_b)/sum_r;
                    if( sum_br > 0.99 && sum_br < 1.20){
                        flabellum.color=PURPLE;
                        continue;
                    }
                    else if(sum_r>sum_b){
                        flabellum.color =  RED;
                    }
                    else if (sum_r < sum_b){
                        flabellum.color = BLUE;
                    }
                    if(centers_.size() != 0) {
                        flabellum.angle = atan2(flabellum.center.y - centers_.at(0).center.y, flabellum.center.x - centers_.at(0).center.x);
                    }
                    flabellum.rotate = rotate;
                    // 保存四个角点
                    cv::Point2f points[4];
                    rotate.points(points);
                    flabellum.points2d.resize(5);
                    flabellum.points2d[0] = points[0];
                    flabellum.points2d[1] = points[1];
                    flabellum.points2d[2] = points[2];
                    flabellum.points2d[3] = points[3];
                    // 按左上左下右下右上的顺序
                    std::sort(flabellum.points2d.begin(), flabellum.points2d.end()-1,[](cv::Point2f &a,cv::Point2f &b){return a.x < b.x;});
                    std::sort(flabellum.points2d.begin(), flabellum.points2d.begin()+1,[](cv::Point2f &a,cv::Point2f &b){return a.y < b.y;});
                    std::sort(flabellum.points2d.begin()+2, flabellum.points2d.end()-1,[](cv::Point2f &a,cv::Point2f &b){return a.y > b.y;});
                    flabellum.points2d[4] = rotate.center;
                    flabellums.emplace_back(flabellum);
                }
            }
        }
        return flabellums;
    }

    // 判断是否是灯条
    bool Detector::isLight(const Light &light)
    {
        // 灯条的长边与短边的比值
        float ratio = light.width / light.length;
        bool ratio_ok = l.min_ratio < ratio && ratio < l.max_ratio;
        bool angle_ok = light.tilt_angle < l.max_angle;
        bool size_ok=(light.width*light.length)>40;
        return ratio_ok && angle_ok&&size_ok;
    }

    // 匹配灯条函数，找到装甲
    std::vector<Armor> Detector::matchLights(const std::vector<Light> &lights, int self_color)
    {
        std::vector<Armor> armors;

        // 遍历所有的灯条组合
        for (int i = 0; i < lights.size(); i++) {
            for (int j = i + 1; j < lights.size(); j++) {
                if (lights[i].color == self_color || lights[j].color == self_color) {
                    continue;
                }

                if (containLight(lights[i], lights[j], lights)) {
                    continue;
                }
                auto armor = Armor(lights[i], lights[j]);
                if (isArmor(armor)) {
                    armors.emplace_back(armor);
                }
            }
        }

        return armors;
    }

    // 判断灯条在boundingRect中的情况下，是否存在其他灯条
    bool Detector::containLight(
            const Light &light_1, const Light &light_2, const std::vector<Light> &lights)
    {
        auto points = std::vector<cv::Point2f>{light_1.top, light_1.bottom, light_2.top, light_2.bottom};
        auto bounding_rect = cv::boundingRect(points);

        for (const auto &test_light: lights) {
            if (test_light.center == light_1.center || test_light.center == light_2.center) {
                continue;
            }

            if (bounding_rect.contains(test_light.top) || bounding_rect.contains(test_light.bottom) ||
                bounding_rect.contains(test_light.center)) {
                return true;
            }
        }

        return false;
    }

    // 判断是否是装甲
    bool Detector::isArmor(Armor &armor)
    {
        Light light_1 = armor.left_light;
        Light light_2 = armor.right_light;
        // 两个灯条的长度比值
        float light_length_ratio = light_1.length < light_2.length ? light_1.length / light_2.length
                                                                   : light_2.length / light_1.length;

        bool light_ratio_ok = light_length_ratio > a.min_light_ratio;



        // 中心灯条的距离（单位：灯条长度）
        float avg_light_length = (light_1.length + light_2.length) / 2;
        float center_distance = cv::norm(light_1.center - light_2.center) / avg_light_length;
        bool center_distance_ok = (a.min_small_center_distance < center_distance &&
                                   center_distance < a.max_small_center_distance) ||
                                  (a.min_large_center_distance < center_distance &&
                                   center_distance < a.max_large_center_distance);

        // 光中心连线的夹角
        cv::Point2f diff = light_1.center - light_2.center;
        float angle = std::abs(std::atan(diff.y / diff.x)) / CV_PI * 180;
        bool angle_ok = angle < a.max_angle;

        // 是否是装甲
        bool is_armor = light_ratio_ok && center_distance_ok && angle_ok;
        armor.armor_type = center_distance > a.min_large_center_distance ? LARGE : SMALL;

        return is_armor;
    }

}// namespace hnurm
